package postgres

import (
	"context"
	"quipus/internal/database/model"
	v1pb "quipus/proto/gen/api/v1"
	apb "quipus/proto/gen/assist"

	_ "github.com/jackc/pgx/v5/stdlib"
	_ "github.com/joho/godotenv/autoload"
	"github.com/pkg/errors"
	"gorm.io/datatypes"
	"gorm.io/gen"
	"gorm.io/gorm/clause"
)

func (p *PGDB) CreateGranary(ctx context.Context, create *model.Granary) error {
	p.query()
	qdo := q.WithContext(ctx)
	err := qdo.Granary.Create(create)
	if err != nil {
		return errors.Wrap(err, "failed to create granary")
	}
	return nil
}

func (p *PGDB) DeleteGranary(ctx context.Context, granayId int32) error {
	p.query()
	g := q.WithContext(ctx).Granary
	_, err := g.Where(q.Granary.ID.Eq(granayId)).Delete()
	if err != nil {
		return errors.Wrap(err, "failed to delete model")
	}
	return nil
}

func (p *PGDB) GetGranary(ctx context.Context, granayId int32) (*model.Granary, error) {
	p.query()
	qdo := q.WithContext(ctx)
	g := qdo.Granary
	return g.Where(q.Granary.ID.Eq(granayId)).First()
}

func (p *PGDB) UpdateGranarySimple(ctx context.Context, update *model.Granary) error {
	p.query()
	qdo := q.WithContext(ctx)
	_, e := qdo.Granary.Updates(update)
	return e
}

func (p *PGDB) UpdateGranary(ctx context.Context, update *apb.UpdateGranary) (*model.Granary, error) {
	p.query()
	qdo := q.WithContext(ctx)
	if update.Id < 0 {
		return nil, errors.New("id is required")
	}

	g := qdo.Granary
	g = g.Where(q.Granary.ID.Eq(update.Id))

	updateValue := &model.Granary{}
	updateValue.ID = update.Id
	if update.UpdatedTs > 0 {
		updateValue.UpdatedTs = update.UpdatedTs
	}

	if update.Name != "" {
		updateValue.Name = update.Name
	}

	if update.Avatar != "" {
		updateValue.AvatarURL = update.Avatar
	}

	if update.Tags != "" {
		updateValue.Tags = update.Tags
	}

	if update.Sources != "" {
		updateValue.Sources = update.Sources
	}

	if len(update.Description) > 0 {
		updateValue.Description = update.Description
	}

	if update.EmbeddingModelId > 0 {
		updateValue.EmbeddingModelID = update.EmbeddingModelId
	}

	if update.Options != "" {
		updateValue.Options = update.Options
	}

	if update.LlmModelId > 0 {
		updateValue.LlmModelID = update.LlmModelId
	}

	_, e := g.Updates(updateValue)

	return updateValue, e
}

func (p *PGDB) ListGranaries(ctx context.Context, find *apb.FindGranary) ([]*model.Granary, int32, error) {
	p.query()
	qdo := q.WithContext(ctx)
	g := qdo.Granary
	if find.Id > 0 {
		g = g.Where(q.LlmModel.ID.Eq(find.Id))
	}

	if len(find.Name) > 0 {
		g = g.Where(q.Granary.Name.Eq(find.Name))
	}

	if find.TenantId > 0 {
		g = g.Where(q.Granary.TenantID.Eq(find.TenantId))
	}

	if len(find.Tags) > 0 {
		jsonCondition := datatypes.JSONQuery("tags")
		for _, tag := range find.Tags {
			g = g.Where(gen.Cond(jsonCondition.Equals(tag))...)
		}
	}

	if len(find.Sources) > 0 {
		jsonCondition := datatypes.JSONQuery("sources")
		for _, s := range find.Sources {
			g = g.Where(gen.Cond(jsonCondition.Equals(s))...)
		}
	}

	limit := find.Limit
	if limit == 0 {
		limit = 10
	}

	ps, c, e := g.FindByPage(int(find.Offset), int(limit))

	return ps, int32(c), e
}

func (p *PGDB) UpsertGranaryResource(ctx context.Context, upsert *model.GranaryResource) error {
	p.query()
	qdo := q.WithContext(ctx)

	err := qdo.GranaryResource.Clauses(clause.OnConflict{
		Columns: []clause.Column{{Name: "granary_id"}, {Name: "resource_id"}},
		DoUpdates: clause.AssignmentColumns([]string{"granary_id", "resource_id", "status",
			"summary", "tags", "retrive_cnt", "created_ts", "updated_ts"}),
	}).Create(upsert)
	if err != nil {
		return err
	}

	return nil
}
func (p *PGDB) DeleteGranaryResource(ctx context.Context, delete *apb.FindGranaryResource) error {
	p.query()
	qdo := q.WithContext(ctx)
	g := qdo.GranaryResource
	if delete.Gid > 0 {
		g = g.Where(q.GranaryResource.GranaryID.Eq(delete.Gid))
	}

	if delete.Rid > 0 {
		g = g.Where(q.GranaryResource.ResourceID.Eq(delete.Rid))
	}
	_, err := g.Delete()
	return err
}

func (p *PGDB) ListGranaryResources(ctx context.Context, find *apb.FindGranaryResource) ([]*model.GranaryResource, error) {
	p.query()
	qdo := q.WithContext(ctx)
	g := qdo.GranaryResource
	if find.Gid > 0 {
		g = g.Where(q.GranaryResource.GranaryID.Eq(find.Gid))
	}

	if find.Rid > 0 {
		g = g.Where(q.GranaryResource.ResourceID.Eq(find.Rid))
	}

	return g.Find()
}

func (p *PGDB) StopGranaryResource(ctx context.Context, find *apb.FindGranaryResource) error {
	p.query()
	qdo := q.WithContext(ctx)
	g := qdo.GranaryResource
	if find.Gid > 0 {
		g = g.Where(q.GranaryResource.GranaryID.Eq(find.Gid))
	}

	if find.Rid > 0 {
		g = g.Where(q.GranaryResource.ResourceID.Eq(find.Rid))
	}

	g = g.Where(q.GranaryResource.Status.Neq(v1pb.GranaryStatus_ENDED.String()))
	g = g.Where(q.GranaryResource.Status.Neq(v1pb.GranaryStatus_STOPPED.String()))
	g = g.Where(q.GranaryResource.Status.Neq(v1pb.GranaryStatus_CREATED.String()))

	updateValue := &model.GranaryResource{Status: v1pb.GranaryStatus_STOPPED.String()}
	_, e := g.Updates(updateValue)
	return e
}
